# 导入tushare
import tushare as ts

# 初始化pro接口
pro = ts.pro_api('69f640cfb31cf07ea489fe1f01a02cf1d1ef52ab9cef755ab0b9aae3')

# 任务一：单只股票数据预处理与指标计算
# ------------------------------------------
# 定义时间范围（示例：获取最近30个交易日）
end_date = "20250331"  # 结束日期
start_date = "20250101"  # 开始日期（确保覆盖足够窗口）

# 拉取单只股票示例数据（以平安银行 000001.SZ 为例）
df = pro.daily(
    ts_code="000001.SZ",
    start_date=start_date,
    end_date=end_date,
    fields=["ts_code", "trade_date", "open", "high", "low", "close"]
)

# 预处理与计算
df = df.sort_values('trade_date')
df['ma5'] = df['close'].rolling(5, min_periods=1).mean()  # 允许最小窗口为1
df['ma10'] = df['close'].rolling(10, min_periods=1).mean()

# RSI计算
delta = df['close'].diff()
gain = delta.where(delta > 0, 0)
loss = -delta.where(delta < 0, 0)
avg_gain = gain.rolling(14, min_periods=1).mean()
avg_loss = loss.rolling(14, min_periods=1).mean()
rs = avg_gain / avg_loss
df['rsi'] = 100 - (100 / (1 + rs))

# 填充缺失值
df.fillna(method='bfill', inplace=True)
print("任务一结果（示例股票指标）：")
print(df[['trade_date', 'close', 'ma5', 'ma10', 'rsi']].tail())

# 任务二：全市场股票金叉点筛选
# ------------------------------------------
# 获取全市场股票列表（剔除ST股、北交所股票）
stock_list = pro.stock_basic(
    exchange='',
    list_status='L',
    fields='ts_code,name,list_date'
)
# 剔除上市不足60日的股票（避免次新股干扰）
stock_list = stock_list[
    pd.to_datetime(stock_list['list_date']) < pd.to_datetime('20240101')
    ]

# 定义输出结果容器
result = []

# 遍历全市场股票
for idx, row in stock_list.iterrows():
    ts_code = row['ts_code']
    try:
        # 获取单只股票日线数据
        df_stock = pro.daily(
            ts_code=ts_code,
            start_date=start_date,
            end_date=end_date,
            fields="ts_code,trade_date,close"
        )
        if len(df_stock) < 10:  # 忽略数据不足的股票
            continue

        # 计算均线
        df_stock = df_stock.sort_values('trade_date')
        df_stock['ma5'] = df_stock['close'].rolling(5, min_periods=1).mean()
        df_stock['ma10'] = df_stock['close'].rolling(10, min_periods=1).mean()
        df_stock.fillna(method='bfill', inplace=True)

        # 判断金叉
        cross_condition = (
                (df_stock['ma5'] > df_stock['ma10']) &
                (df_stock['ma5'].shift(1) <= df_stock['ma10'].shift(1))
        )
        cross_dates = df_stock[cross_condition]

        if not cross_dates.empty:
            for _, cross_row in cross_dates.iterrows():
                result.append({
                    'ts_code': ts_code,
                    'name': row['name'],
                    'trade_date': cross_row['trade_date'],
                    'close_price': cross_row['close']
                })
    except Exception as e:
        print(f"处理股票 {ts_code} 时出错: {str(e)}")

# 转换为DataFrame并输出
result_df = pd.DataFrame(result)
print("\n任务二结果（全市场金叉股票）：")
print(result_df.head())